(*astToIR.ml*)
(* 引入 AST 和 IR 类型 *)
open Ast
open Ir
open Optimazation

(* Map<String, operand> *)
module Env = Map.Make (String)

module EnvStack = struct
  type t = operand Env.t list

  let empty : t = [ Env.empty ]
  let enter (stk : t) : t = Env.empty :: stk

  let exit = function
    | _ :: tl -> tl
    | [] -> failwith "EnvStack.exit: empty stack"
  ;;

  let add name v = function
    | top :: tl -> Env.add name v top :: tl
    | [] -> failwith "EnvStack.add: empty stack"
  ;;

  let rec lookup name = function
    | [] -> failwith ("unbound variable: " ^ name)
    | m :: ms ->
      (match Env.find_opt name m with
       | Some v -> v
       | None -> lookup name ms)
  ;;
end

(* 翻译上下文：包含当前 env 及循环标签，用于实现 break/continue *)
type context =
  { func_name : string (* 当前函数名 *)
  ; env_stack : EnvStack.t ref
  ; break_lbl : string option (* break 跳转目标 *)
  ; continue_lbl : string option (* continue 跳转目标 *)
  }

module StringSet = Set.Make (String)

(* 1. 副作用判断：只把 Call/Return/Break/Continue 视为副作用 *)
let rec has_side_effect_expr = function
  | Ast.Call _ -> true
  | Ast.Unop (_, e) -> has_side_effect_expr e
  | Ast.Binop (_, l, r) -> has_side_effect_expr l || has_side_effect_expr r
  | _ -> false
;;

let rec has_side_effect_stmt = function
  | Ast.ExprStmt e -> has_side_effect_expr e
  | Ast.Return _ -> true
  | Ast.Break | Ast.Continue -> true
  | Ast.If (_, t, f) ->
    has_side_effect_stmt t
    ||
      (match f with
      | Some f -> has_side_effect_stmt f
      | None -> false)
  | Ast.While (_, b) -> has_side_effect_stmt b
  | Ast.Block ss -> List.exists has_side_effect_stmt ss
  | _ -> false (* 不把 Assign/Decl 视为副作用 *)
;;

(* 2. 收集变量写入 *)
let rec vars_written_stmt = function
  | Ast.Assign (x, _) -> StringSet.singleton x
  | Ast.Decl (x, _) -> StringSet.singleton x
  | Ast.If (_, t, f) ->
    let w1 = vars_written_stmt t in
    let w2 =
      match f with
      | Some f -> vars_written_stmt f
      | None -> StringSet.empty
    in
    StringSet.union w1 w2
  | Ast.While (_, b) -> vars_written_stmt b
  | Ast.Block ss ->
    List.fold_left
      (fun acc s -> StringSet.union acc (vars_written_stmt s))
      StringSet.empty
      ss
  | _ -> StringSet.empty
;;

(* 3. 收集变量读取，用于更新 seen *)
let rec vars_read_expr = function
  | Ast.ID x -> StringSet.singleton x
  | Ast.Unop (_, e) -> vars_read_expr e
  | Ast.Binop (_, l, r) -> StringSet.union (vars_read_expr l) (vars_read_expr r)
  | Ast.Call (_, args) ->
    List.fold_left
      (fun acc e -> StringSet.union acc (vars_read_expr e))
      StringSet.empty
      args
  | _ -> StringSet.empty
;;

let rec vars_read_stmt = function
  | Ast.ExprStmt e -> vars_read_expr e
  | Ast.Return (Some e) -> vars_read_expr e
  | Ast.Return None -> StringSet.empty
  | Ast.Decl (_, Some e) -> vars_read_expr e
  | Ast.Assign (_, e) -> vars_read_expr e
  | Ast.If (c, t, f) ->
    StringSet.union
      (vars_read_expr c)
      (StringSet.union
         (vars_read_stmt t)
         (match f with
          | Some f -> vars_read_stmt f
          | None -> StringSet.empty))
  | Ast.While (c, b) -> StringSet.union (vars_read_expr c) (vars_read_stmt b)
  | Ast.Block ss ->
    List.fold_left
      (fun acc s -> StringSet.union acc (vars_read_stmt s))
      StringSet.empty
      ss
  | _ -> StringSet.empty
;;

(* 4. 删除无用 while *)
(* 4. 删除无用 while：逆序遍历，确保基于后续使用消除 *)
let rec remove_dead_while stmts =
  (* 对每个 While，查看其写入的变量在后续语句中是否被读取 *)
  let rec go acc = function
    | [] -> List.rev acc
    | stmt :: rest ->
      let keep =
        match stmt with
        | Ast.While (_, body) ->
          let se = has_side_effect_stmt body in
          let writes = vars_written_stmt body in
          (* 计算后续语句中所有读取的变量 *)
          let future_reads =
            List.fold_left
              (fun s st -> StringSet.union s (vars_read_stmt st))
              StringSet.empty
              rest
          in
          (* 如果无副作用且写入的变量都未被后续读取，则删除 *)
          not ((not se) && StringSet.is_empty (StringSet.inter writes future_reads))
        | Ast.Block _ -> true
        | _ -> true
      in
      let acc' =
        if keep
        then (
          let stmt' =
            match stmt with
            | Ast.Block ss -> Ast.Block (remove_dead_while ss)
            | other -> other
          in
          stmt' :: acc)
        else acc
      in
      go acc' rest
  in
  go [] stmts
;;

let rec contains_while = function
  | [] -> false
  | Ast.While _ :: _ -> true
  | Ast.Block ss :: rest -> contains_while ss || contains_while rest
  | _ :: rest -> contains_while rest
;;

(* 只有当函数体里真正包含 while 时，才调用 remove_dead_while *)
let preprocess_ast (cu : Ast.comp_unit) : Ast.comp_unit =
  List.map
    (fun f ->
       if contains_while f.Ast.body
       then { f with Ast.body = remove_dead_while f.Ast.body }
       else f (* 不含 while 的函数保持原 AST，不做任何修改 *))
    cu
;;

let rec is_const_expr = function
  | Number _ -> true
  | Binop (_, e1, e2) -> is_const_expr e1 && is_const_expr e2
  | _ -> false
;;

let rec uses_var_in_expr name = function
  | ID id -> id = name
  | Binop (_, e1, e2) -> uses_var_in_expr name e1 || uses_var_in_expr name e2
  | Unop (_, e) -> uses_var_in_expr name e
  | Call (_, args) -> List.exists (uses_var_in_expr name) args
  | _ -> false
;;

let rec uses_var_in_stmt name stmt =
  match stmt with
  | Ast.Assign (_, e) -> uses_var_in_expr name e
  | Decl (_, Some e) -> uses_var_in_expr name e
  | ExprStmt e -> uses_var_in_expr name e
  | If (cond, s1, s2_opt) ->
    uses_var_in_expr name cond
    || uses_var_in_stmt name s1
    ||
      (match s2_opt with
      | Some s2 -> uses_var_in_stmt name s2
      | None -> false)
  | While (cond, s) -> uses_var_in_expr name cond || uses_var_in_stmt name s
  | Block ss -> List.exists (uses_var_in_stmt name) ss
  | _ -> false
;;

(* 尝试消除嵌套累加型 while 循环 *)
let rec elim_loop_stmt (stmt : stmt) : stmt =
  match stmt with
  (* 只支持 < 不支持 <= *)
  | While (Binop (Less, ID idx, Number n), Block body) ->
    (* 尝试匹配固定结构 *)
    let match_inner_loop stmts =
      match stmts with
      | Decl (k_name, Some (Number 0)) :: assigns_before
        when List.exists
               (function
                 | Ast.Assign (_, Number 0) -> true
                 | _ -> false)
               assigns_before ->
        (* 找出累加器初始化的 assign 语句 *)
        let acc_inits, rest =
          List.partition
            (function
              | Ast.Assign (_, Number 0) -> true
              | _ -> false)
            assigns_before
        in
        let acc_names =
          List.filter_map
            (function
              | Ast.Assign (name, Number 0) -> Some name
              | _ -> None)
            acc_inits
        in
        (* 找出内层循环 *)
        (match rest with
         | While (Binop (Less, ID k_id, Number m), Block inner_body) :: tail
           when k_id = k_name ->
           (* 检查 inner_body 是否符合 acc = acc + expr 的形式 *)
           let valid_acc_exprs =
             List.filter_map
               (function
                 | Ast.Assign (acc, Binop (Add, ID acc2, expr))
                   when acc = acc2 && List.mem acc acc_names -> Some (acc, expr)
                 | _ -> None)
               inner_body
           in
           let valid =
             List.length valid_acc_exprs = List.length acc_names
             && List.for_all
                  (fun acc -> List.mem acc acc_names)
                  (List.map fst valid_acc_exprs)
           in
           if valid then Some (k_name, m, valid_acc_exprs, tail) else None
         | _ -> None)
      | _ -> None
    in
    (match match_inner_loop body with
     | Some (k_var, m, acc_exprs, tail_after_loop) ->
       let _ = List.map fst acc_exprs in
       (* 替换所有 acc = acc + expr 为 acc = expr * m *)
       let new_accs =
         List.map
           (fun (acc, expr) -> Ast.Assign (acc, Binop (Mul, expr, Number m)))
           acc_exprs
       in
       (* 如果后续用到了 k_var，还要补 Decl(k, 0) *)
       let k_decl =
         if List.exists (uses_var_in_stmt k_var) tail_after_loop
         then [ Decl (k_var, Some (Number 0)) ]
         else []
       in
       (* 构造新循环体 *)
       let new_body =
         Block (k_decl @ new_accs @ List.map elim_loop_stmt tail_after_loop)
       in
       While (Binop (Less, ID idx, Number n), new_body)
     | None ->
       (* 非匹配模式，递归优化内部语句 *)
       While (Binop (Less, ID idx, Number n), Block (List.map elim_loop_stmt body)))
  | While (cond, Block body) -> While (cond, Block (List.map elim_loop_stmt body))
  | Block stmts -> Block (List.map elim_loop_stmt stmts)
  | If (cond, t_branch, f_branch) ->
    If (cond, elim_loop_stmt t_branch, Option.map elim_loop_stmt f_branch)
  | _ -> stmt
;;

let rec remove_trivial_self_inc_loop stmt =
  match stmt with
  | While (Binop (Less, ID var, Number _),
           Block [Assign (var2, Binop (Add, ID var3, Number 1))])
    when var = var2 && var = var3 ->
      (* 无效果自增循环，删除 *)
      Block []

  | Block stmts ->
      Block (stmts |> List.map remove_trivial_self_inc_loop |> List.filter (function Block [] -> false | _ -> true))

  | If (cond, s1, s2_opt) ->
      If (
        cond,
        remove_trivial_self_inc_loop s1,
        Option.map remove_trivial_self_inc_loop s2_opt
      )

  | While (cond, Block body) ->
      While (cond, Block (List.map remove_trivial_self_inc_loop body))

  | other -> other


let elim_loop_func (f : func_def) : func_def =
    let new_body =
      f.body
      |> List.map elim_loop_stmt
      |> List.map remove_trivial_self_inc_loop
    in
    { f with body = new_body }
;;

let loop_elim_ast (cu : comp_unit) : comp_unit = List.map elim_loop_func cu

module LabelMap = Map.Make (String)

(* 临时寄存器与标签生成器 *)
let temp_id = ref 0

let fresh_temp () =
  let id = !temp_id in
  incr temp_id;
  Reg ("t" ^ string_of_int id)
;;

let name_id = ref 0

let fresh_name base =
  let id = !name_id in
  incr name_id;
  base ^ "_" ^ string_of_int id
;;

let label_id = ref 0
let ir_label_id = ref 0

let fresh_label () =
  let id = !label_id in
  incr label_id;
  "L" ^ string_of_int id
;;

let fresh_IRlabel (label_map : string LabelMap.t) (l : param) : string * string LabelMap.t
  =
  (* 先查找 l 是否已经分配了 label，如果有直接返回，否则分配新 label，并写回 map *)
  match LabelMap.find_opt l label_map with
  | Some lbl -> lbl, label_map
  | None ->
    let id = !ir_label_id in
    incr ir_label_id;
    let lbl = "LABEL" ^ string_of_int id in
    let label_map' = LabelMap.add l lbl label_map in
    lbl, label_map'
;;

(* 操作符到字符串映射 *)
let string_of_binop = function
  | Add -> "+"
  | Sub -> "-"
  | Mul -> "*"
  | Div -> "/"
  | Mod -> "%"
  | Eq -> "=="
  | Neq -> "!="
  | Less -> "<"
  | Leq -> "<="
  | Greater -> ">"
  | Geq -> ">="
  | Land -> "&&"
  | Lor -> "||"
;;

let string_of_unop = function
  | Not -> "!"
  | Plus -> "+"
  | Minus -> "-"
;;

(* stmt_to_res 用于处理 return 终止：Normal/Returned 两种结果 *)
type stmt_res =
  | Normal of ir_inst list
  | Returned of ir_inst list

(* 将 stmt_res 展平为代码列表 *)
let flatten = function
  | Normal code | Returned code -> code
;;

(*判断块是否return*)
let always_returns (res : stmt_res) : bool =
  match res with
  | Returned _ -> true
  | Normal _ -> false
;;

(* 检查代码段最后一条是否是 Goto 指定标签或 Return *)
let ends_with_jump_or_return insts =
  match List.rev insts with
  | Goto _ :: _ -> true
  | Ret _ :: _ -> true
  | _ -> false
;;

(* 表达式转换：返回目标寄存器和 IR 指令列表 *)
let rec expr_to_ir (ctx : context) (e : expr) : operand * ir_inst list =
  match e with
  | Number n -> Imm n, []
  | ID name ->
    let operand = EnvStack.lookup name !(ctx.env_stack) in
    operand, []
  | Unop (op, e1) ->
    let operand, code = expr_to_ir ctx e1 in
    (match operand with
     | Imm n ->
       let folded =
         match op with
         | Plus -> Imm n (* +n = n *)
         | Minus -> Imm (-n) (* -n *)
         | Not -> Imm (if n = 0 then 1 else 0)
         (* !n *)
       in
       folded, code
     | _ ->
       let res = fresh_temp () in
       res, code @ [ Unop (string_of_unop op, res, operand) ])
  | Binop (op, e1, e2) ->
    let lhs, c1 = expr_to_ir ctx e1 in
    let rhs, c2 = expr_to_ir ctx e2 in
    (match op with
     | Land | Lor ->
       let dst = fresh_temp () in
       dst, c1 @ c2 @ [ Binop (string_of_binop op, dst, lhs, rhs) ]
     | _ ->
       (match lhs, rhs with
        | Imm a, Imm b ->
          let folded =
            match op with
            | Add -> Imm (a + b)
            | Sub -> Imm (a - b)
            | Mul -> Imm (a * b)
            | Div -> if b <> 0 then Imm (a / b) else Imm 9999
            | Mod -> Imm (a mod b)
            | Eq -> Imm (if a = b then 1 else 0)
            | Neq -> Imm (if a <> b then 1 else 0)
            | Less -> Imm (if a < b then 1 else 0)
            | Leq -> Imm (if a <= b then 1 else 0)
            | Greater -> Imm (if a > b then 1 else 0)
            | Geq -> Imm (if a >= b then 1 else 0)
            | _ -> failwith "Never touched"
          in
          folded, c1 @ c2
        | _ ->
          let dst = fresh_temp () in
          dst, c1 @ c2 @ [ Binop (string_of_binop op, dst, lhs, rhs) ]))
  | Call (f, args) ->
    (* 参数顺序按出现顺序计算 *)
    let codes, oprs =
      List.fold_left
        (fun (acc_code, acc_opr) arg ->
           let opr, code = expr_to_ir ctx arg in
           acc_code @ code, acc_opr @ [ opr ])
        ([], [])
        args
    in
    let ret = fresh_temp () in
    ret, codes @ [ Call (ret, f, oprs) ]
;;

let rec normalize_expr = function
  | Ast.Unop (Not, Unop (Not, e)) -> normalize_expr e (* !!e => e *)
  | Unop (Not, Binop (Land, a, b)) ->
    (* !(a && b) => !a || !b *)
    normalize_expr (Binop (Lor, Unop (Not, a), Unop (Not, b)))
  | Unop (Not, Binop (Lor, a, b)) ->
    (* !(a || b) => !a && !b *)
    normalize_expr (Binop (Land, Unop (Not, a), Unop (Not, b)))
  | Unop (Not, Binop (op, a, b)) ->
    (* 关系运算反转：!(x > y) => x <= y *)
    let neg =
      match op with
      | Eq -> Neq
      | Neq -> Eq
      | Less -> Geq
      | Leq -> Greater
      | Greater -> Leq
      | Geq -> Less
      | _ -> failwith "unsupported negation of this binary op"
    in
    Ast.Binop (neg, normalize_expr a, normalize_expr b)
  | Binop (op, a, b) -> Binop (op, normalize_expr a, normalize_expr b)
  | Unop (op, a) -> Unop (op, normalize_expr a)
  | Call (f, args) -> Call (f, List.map normalize_expr args)
  | e -> e
;;

(* 将 if(cond, then_b, else_b) 里的 cond 展开 *)
let rec desugar_stmt = function
  | If (cond, then_b, Some else_b) ->
    let cond = normalize_expr cond in
    let ands = split_and cond in
    if List.length ands > 1
    then (
      (* a && b && c => if a then if b then ... else else_b *)
      let rec nest_and = function
        | [ x ] -> If (x, then_b, Some else_b)
        | hd :: tl -> If (hd, Block [ nest_and tl ], Some else_b)
        | [] -> Block []
      in
      nest_and ands |> desugar_stmt)
    else (
      let ors = split_or cond in
      if List.length ors > 1
      then (
        (* a || b || c => if a then then_b else if b ... *)
        let rec nest_or = function
          | [ x ] -> If (x, then_b, Some else_b)
          | hd :: tl -> If (hd, then_b, Some (nest_or tl))
          | [] -> Block []
        in
        nest_or ors |> desugar_stmt)
      else If (cond, desugar_stmt then_b, Some (desugar_stmt else_b)))
  | If (cond, then_b, None) ->
    let cond = normalize_expr cond in
    let ands = split_and cond in
    if List.length ands > 1
    then (
      let rec nest_and = function
        | [ x ] -> If (x, then_b, None)
        | hd :: tl -> If (hd, Block [ nest_and tl ], None)
        | [] -> Block []
      in
      nest_and ands |> desugar_stmt)
    else (
      let ors = split_or cond in
      if List.length ors > 1
      then (
        let rec nest_or = function
          | [ x ] -> If (x, then_b, None)
          | hd :: tl -> If (hd, then_b, Some (nest_or tl))
          | [] -> Block []
        in
        nest_or ors |> desugar_stmt)
      else If (cond, desugar_stmt then_b, None))
  | While (cond, body) -> While (normalize_expr cond, desugar_stmt body)
  | Block ss -> Block (List.map desugar_stmt ss)
  | other -> other
;;

(* 语句翻译，返回 Normal/Returned，支持块作用域、break/continue、return 提前终止 *)
let rec stmt_to_res (ctx : context) (in_tail : bool) (s : stmt) : stmt_res =
  match s with
  | Empty -> Normal []
  | ExprStmt e ->
    let _, code = expr_to_ir ctx e in
    Normal code
  | Decl (x, None) ->
    let new_name = fresh_name x in
    ctx.env_stack := EnvStack.add x (Var new_name) !(ctx.env_stack);
    Normal []
  | Decl (x, Some e) ->
    let v, c = expr_to_ir ctx e in
    let new_name = fresh_name x in
    ctx.env_stack := EnvStack.add x (Var new_name) !(ctx.env_stack);
    Normal (c @ [ Assign (Var new_name, v) ])
  | Assign (x, e) ->
    let v, c = expr_to_ir ctx e in
    let var = EnvStack.lookup x !(ctx.env_stack) in
    Normal (c @ [ Assign (var, v) ])
  | Return None -> if in_tail then Returned [] else Returned [ Ret None ]
  | Return (Some e) ->
    (match e with
     (* 只有对自身的调用才做尾调用消除 *)
     | Call (f, args) when f = ctx.func_name ->
       let arg_codes, arg_oprs =
         List.fold_left
           (fun (cc, oo) arg ->
              let o, c = expr_to_ir ctx arg in
              cc @ c, oo @ [ o ])
           ([], [])
           args
       in
       Returned (arg_codes @ [ TailCall (f, arg_oprs) ])
     (* 其他一律当普通表达式处 理 *)
     | _ ->
       let r, code = expr_to_ir ctx e in
       Returned (code @ [ Ret (Some r) ]))
  | If (cond, tstmt, Some fstmt) ->
    let cnd, cc = expr_to_ir ctx cond in
    let lthen = fresh_label ()
    and lelse = fresh_label ()
    and lend = fresh_label () in
    let then_res = stmt_to_res ctx in_tail tstmt
    and else_res = stmt_to_res ctx in_tail fstmt in
    let raw_then = flatten then_res in
    let then_code =
      if ends_with_jump_or_return raw_then then raw_then else raw_then @ [ Goto lend ]
    in
    let raw_else = flatten else_res in
    let else_code =
      if ends_with_jump_or_return raw_else then raw_else else raw_else @ [ Goto lend ]
    in
    let code =
      cc
      @ [ IfGoto (cnd, lthen); Goto lelse ]
      @ [ Label lthen ]
      @ then_code
      @ [ Label lelse ]
      @ else_code
      @ [ Label lend ]
    in
    if always_returns then_res && always_returns else_res
    then Returned code
    else Normal code
  | If (cond, tstmt, None) ->
    let cnd, cc = expr_to_ir ctx cond in
    let lthen = fresh_label ()
    and lskip = fresh_label () in
    let then_res = stmt_to_res ctx in_tail tstmt in
    let then_code = flatten then_res in
    let code =
      cc
      @ [ IfGoto (cnd, lthen); Goto lskip ]
      @ [ Label lthen ]
      @ then_code
      @ [ Label lskip ]
    in
    Normal code
  | While (cond, body) ->
    (* 循环标签 *)
    let lcond = fresh_label ()
    and lbody = fresh_label ()
    and lend = fresh_label () in
    let ctx_loop = { ctx with break_lbl = Some lend; continue_lbl = Some lcond } in
    let cnd, ccode = expr_to_ir ctx_loop cond in
    let body_res = stmt_to_res ctx_loop false body in
    let bcode = flatten body_res in
    let code =
      [ Goto lcond; Label lcond ]
      @ ccode
      @ [ IfGoto (cnd, lbody); Goto lend ]
      @ [ Label lbody ]
      @ bcode
      @ [ Goto lcond; Label lend ]
    in
    (* 无法从循环体中直接 return：若想支持可在 body_res 捕获 *)
    Normal code
  | Break ->
    (match ctx.break_lbl with
     | Some lbl -> Normal [ Goto lbl ]
     | None -> failwith "break used outside loop")
  | Continue ->
    (match ctx.continue_lbl with
     | Some lbl -> Normal [ Goto lbl ]
     | None -> failwith "continue used outside loop")
  | Block stmts ->
    ctx.env_stack := EnvStack.enter !(ctx.env_stack);
    let rec loop acc = function
      | [] -> Normal acc
      | [ last ] ->
        (match stmt_to_res ctx in_tail last with
         | Returned c -> Returned (acc @ c)
         | Normal c -> Normal (acc @ c))
      | hd :: tl ->
        (match stmt_to_res ctx false hd with
         | Returned c -> Returned (acc @ c)
         | Normal c -> loop (acc @ c) tl)
    in
    let res = loop [] stmts in
    ctx.env_stack := EnvStack.exit !(ctx.env_stack);
    res
;;

(* 函数转换 *)
let func_to_ir (f : func_def) : ir_func =
  (* 1. 预处理 AST *)
  let desugared_body =
    match desugar_stmt (Block f.body) with
    | Block ss -> ss
    | _ -> f.body
  in
  let f' = { f with body = desugared_body } in
  (* 2. 初始化环境和上下文 *)
  let init_env = List.fold_left (fun m p -> Env.add p (Var p) m) Env.empty f'.params in
  let ctx0 =
    { func_name = f'.func_name
    ; env_stack = ref [ init_env ]
    ; break_lbl = None
    ; continue_lbl = None
    }
  in
  (* 3. 生成完整 IR，不当全局尾调用上下文 *)
  let raw_res = stmt_to_res ctx0 false (Block f'.body) in
  let raw_code = flatten raw_res in
  (* 4. 获取参数对应的 operand 列表，用于尾调用展开 *)
  let param_ops = List.map (fun p -> Var p) f'.params in
  (* 5. 插入循环入口标签 *)
  let entry_lbl = fresh_label () in
  let code_with_entry = Label entry_lbl :: raw_code in
  (* 6. 展开尾调用为参数赋值 + Goto *)
  let final_code =
    List.concat_map
      (fun inst ->
         match inst with
         | TailCall (fname, args) when fname = f'.func_name ->
           List.map2 (fun param_op arg_op -> Assign (param_op, arg_op)) param_ops args
           @ [ Goto entry_lbl ]
         | _ -> [ inst ])
      code_with_entry
  in
  { name = f'.func_name; args = f'.params; body = final_code }
;;

(* 线性IR -> 过程块IR *)
let partition_blocks (insts : ir_inst list) : ir_block list =
  let rec split acc curr label label_map insts =
    match insts with
    | [] -> List.rev acc (*| _ -> failwith "Basic block must end with a terminator")*)
    | Label l :: rest ->
      (* 当前块结束，开启新块 *)
      (match curr with
       | [] ->
         let next_label, label_map' = fresh_IRlabel label_map l in
         split acc [ Label l ] next_label label_map' rest
       | _ ->
         let next_label, label_map' = fresh_IRlabel label_map l in
         let blk =
           { label
           ; insts = List.rev curr
           ; terminator = TermSeq next_label
           ; preds = []
           ; succs = []
           ; live_in = OperandSet.empty
           ; live_out = OperandSet.empty
           }
         in
         let acc' = blk :: acc in
         split acc' [ Label l ] next_label label_map' rest)
    | Goto l :: rest ->
      let goto_label, label_map' = fresh_IRlabel label_map l in
      (* 刷新一个无意义的 blk, 确保编程者不会出现的 label *)
      let next_label, label_map'' =
        fresh_IRlabel label_map' ("__blk" ^ string_of_int !ir_label_id)
      in
      let blk =
        { label
        ; insts = List.rev (Goto l :: curr)
        ; terminator = TermGoto goto_label
        ; preds = []
        ; succs = []
        ; live_in = OperandSet.empty
        ; live_out = OperandSet.empty
        }
      in
      split (blk :: acc) [] next_label label_map'' rest
    | IfGoto (cond, l) :: rest ->
      let then_label, label_map' = fresh_IRlabel label_map l in
      let else_label, label_map'' =
        fresh_IRlabel label_map' ("__else" ^ string_of_int !ir_label_id)
      in
      let blk =
        { label
        ; insts = List.rev (IfGoto (cond, l) :: curr)
        ; terminator = TermIf (cond, then_label, else_label)
        ; preds = []
        ; succs = []
        ; live_in = OperandSet.empty
        ; live_out = OperandSet.empty
        }
      in
      split (blk :: acc) [] else_label label_map'' rest
    | Ret op :: rest ->
      let next_label, label_map' =
        fresh_IRlabel label_map ("__ret" ^ string_of_int !ir_label_id)
      in
      let blk =
        { label
        ; insts = List.rev (Ret op :: curr)
        ; terminator = TermRet op
        ; preds = []
        ; succs = []
        ; live_in = OperandSet.empty
        ; live_out = OperandSet.empty
        }
      in
      split (blk :: acc) [] next_label label_map' rest
    | inst :: rest -> split acc (inst :: curr) label label_map rest
  in
  (* 确保用户不使用 entry 标签 *)
  let entry_label, label_map = fresh_IRlabel LabelMap.empty "entry" in
  split [] [] entry_label label_map insts
;;

(* let show_operand = function
  | Var x -> x
  | Reg r -> r
  | Imm n -> string_of_int n

(* 打印终结符 *)
let show_terminator = function
  | TermGoto lbl -> Printf.sprintf "goto %s" lbl
  | TermIf (cond, t, f) -> 
      Printf.sprintf "if %s then %s else %s" 
        (show_operand cond) t f
  | TermRet op -> 
      Printf.sprintf "ret %s" 
        (match op with Some v -> show_operand v | None -> "")
  | TermSeq lbl -> Printf.sprintf "seq %s" lbl

(* 打印单条IR指令 *)
let show_ir_inst = function
  | Binop (op, dst, lhs, rhs) ->
      Printf.sprintf "%s = %s %s %s" 
        (show_operand dst) (show_operand lhs) op (show_operand rhs)
  | Unop (op, dst, src) ->
      Printf.sprintf "%s = %s %s" 
        (show_operand dst) op (show_operand src)
  | Assign (dst, src) ->
      Printf.sprintf "%s = %s" 
        (show_operand dst) (show_operand src)
  | Call (dst, f, args) ->
      Printf.sprintf "%s = call %s(%s)" 
        (show_operand dst) f (List.map show_operand args |> String.concat ", ")
  | Ret op -> 
      Printf.sprintf "ret %s" 
        (match op with Some v -> show_operand v | None -> "")
  | Goto lbl -> 
      Printf.sprintf "goto %s" lbl
  | IfGoto (cond, lbl) -> 
      Printf.sprintf "if %s goto %s" 
        (show_operand cond) lbl
  | Label lbl -> 
      Printf.sprintf "%s:" lbl
  | _ -> "unknown_instruction"

打印操作数 *)
let has_side_effect inst =
  match inst with
  | Call _ | Store _ | Ret _ -> true
  | Goto _ | IfGoto _ | Label _ -> true
  (* | Assign (Var _, Imm 0) -> true *)
  | _ -> false
;;

let dead_code_elimination blocks (print_liveness : bool) =
  (* 1. 先做活跃变量分析 *)
  liveness_analysis blocks print_liveness;
  (* 2. 对每个基本块做 DCE *)
  List.map
    (fun blk ->
       let live = ref blk.live_out in
       (* 倒序处理指令 *)
       let new_insts =
         List.fold_right
           (fun inst acc ->
              let def, use = def_use_inst inst in
              let must_keep = has_side_effect inst in
              let def_is_live =
                (not (OperandSet.is_empty def))
                && OperandSet.exists (fun v -> OperandSet.mem v !live) def
              in
              if must_keep || def_is_live
              then (
                (* 如果保留，就更新 live 集合 *)
                live := OperandSet.union use (OperandSet.diff !live def);
                inst :: acc)
              else acc)
           blk.insts
           []
       in
       (* 返回更新后的基本块 *)
       { blk with insts = new_insts })
    blocks
;;

(* 先只二元运算做 CSE *)
module ExprMap = Map.Make (struct
    type t = string * operand * operand (* op, lhs, rhs *)

    let compare = compare
  end)

let cse_block (blk : ir_block) : ir_block =
  let expr_table = ref ExprMap.empty in
  let new_insts =
    List.fold_left
      (fun acc inst ->
         match inst with
         | Binop (op, dst, lhs, rhs) ->
           let key = op, lhs, rhs in
           (match ExprMap.find_opt key !expr_table with
            | Some prev ->
              (* 替换成 Assign *)
              Assign (dst, prev) :: acc
            | None ->
              (* 插入新表达式 *)
              expr_table := ExprMap.add key dst !expr_table;
              inst :: acc)
         | _ ->
           (* 遇到 Store/Call/Ret 等副作用或不可预测的语义，应清空 expr_table *)
           let kills_all =
             match inst with
             | Store _ | Call _ | TailCall _ | Ret _ -> true
             | _ -> false
           in
           if kills_all then expr_table := ExprMap.empty;
           inst :: acc)
      []
      blk.insts
    |> List.rev
  in
  { blk with insts = new_insts }
;;

module OperandMap = Map.Make (struct
    type t = operand

    let compare = compare
  end)

let rec resolve_copy env op =
  match op with
  | Var _ | Reg _ ->
    (match OperandMap.find_opt op env with
     | Some v when v <> op -> resolve_copy env v
     | _ -> op)
  | _ -> op
;;

let copy_propagation_block (blk : ir_block) : ir_block =
  let copy_env = ref OperandMap.empty in
  let propagate_op op = resolve_copy !copy_env op in
  let new_insts =
    List.filter_map
      (fun inst ->
         match inst with
         | Assign (dst, src) ->
           let src' = propagate_op src in
           (* 如果 dst = src 形成 a = a, 可以删掉 *)
           if dst = src'
           then None
           else (
             copy_env := OperandMap.add dst src' !copy_env;
             Some (Assign (dst, src')))
         | Binop (op, dst, lhs, rhs) ->
           let lhs' = propagate_op lhs in
           let rhs' = propagate_op rhs in
           copy_env := OperandMap.remove dst !copy_env;
           Some (Binop (op, dst, lhs', rhs'))
         | Unop (op, dst, src) ->
           let src' = propagate_op src in
           copy_env := OperandMap.remove dst !copy_env;
           Some (Unop (op, dst, src'))
         | Load (dst, src) ->
           let src' = propagate_op src in
           copy_env := OperandMap.remove dst !copy_env;
           Some (Load (dst, src'))
         | Store (dst, src) ->
           let dst' = propagate_op dst in
           let src' = propagate_op src in
           Some (Store (dst', src'))
         | Ret (Some src) ->
           let src' = propagate_op src in
           Some (Ret (Some src'))
         | Call (dst, fname, args) ->
           let args' = List.map propagate_op args in
           (* let dst_id = match dst with Var v | Reg v -> v in *)
           copy_env := OperandMap.remove dst !copy_env;
           Some (Call (dst, fname, args'))
         | TailCall (fname, args) ->
           let args' = List.map propagate_op args in
           Some (TailCall (fname, args'))
         | Goto _ | IfGoto _ | Label _ | Ret None -> Some inst)
      blk.insts
  in
  { blk with insts = new_insts }
;;

module StringMap = Map.Make (String)

let dce2_block (blk : ir_block) : unit =
  (* 1. 构建 def map: Reg/Var -> inst *)
  let def_map =
    List.fold_left
      (fun acc inst ->
         match inst with
         | Binop (_, dst, _, _)
         | Unop (_, dst, _)
         | Assign (dst, _)
         | Load (dst, _)
         | Call (dst, _, _) ->
           (match dst with
            | Reg r | Var r -> StringMap.add r inst acc
            | _ -> acc)
         | _ -> acc)
      StringMap.empty
      blk.insts
  in
  (* 2. 获取终结符使用的根变量（如 return Reg t10） *)
  (* 提取 terminator 中所有使用到的变量 *)
  let term_vars =
    match blk.terminator with
    | TermRet (Some op) | TermIf (op, _, _) ->
      (match op with
       | Reg r | Var r -> StringSet.singleton r
       | _ -> StringSet.empty)
    | _ -> StringSet.empty
  in
  let roots = term_vars in
  (* 3. worklist 算法向上追踪依赖 *)
  let rec collect acc worklist =
    match worklist with
    | [] -> acc
    | v :: rest when StringSet.mem v acc -> collect acc rest
    | v :: rest ->
      let acc = StringSet.add v acc in
      let new_vars =
        match StringMap.find_opt v def_map with
        | Some inst ->
          (* 提取 inst 中使用的变量 *)
          let use_vars =
            match inst with
            | Binop (_, _, lhs, rhs) | Store (lhs, rhs) -> [ lhs; rhs ]
            | Unop (_, _, src) | Load (_, src) | Assign (_, src) -> [ src ]
            | Call (_, _, args) -> args
            | _ -> []
          in
          use_vars
          |> List.filter_map (function
            | Reg r | Var r -> Some r
            | _ -> None)
        | None -> []
      in
      collect acc (new_vars @ rest)
  in
  let useful_vars = collect StringSet.empty (StringSet.elements roots) in
  (* 4. 只保留用到的 insts *)
  blk.insts
  <- List.filter
       (fun inst ->
          match inst with
          | Binop (_, dst, _, _)
          | Unop (_, dst, _)
          | Assign (dst, _)
          | Load (dst, _)
          | Call (dst, _, _) ->
            (match dst with
             | Reg r | Var r -> StringSet.mem r useful_vars
             | _ -> false)
          | _ -> true (* Label, Goto 等不删除 *))
       blk.insts
;;

(* 优化版本的 ir 控制块 *)
(* func_to_ir_o: 优化版，支持尾调用消除后构建 CFG *)
let func_to_ir_o (f : func_def) (print_liveness : bool) : ir_func_o =
  (* 1. 初始化环境 *)
  let desugared_body =
    match desugar_stmt (Block f.body) with
    | Block ss -> ss
    | _ -> f.body
  in
  let f' = { f with body = desugared_body } in
  let init_env = List.fold_left (fun m p -> Env.add p (Var p) m) Env.empty f'.params in
  let ctx0 =
    { func_name = f'.func_name
    ; env_stack = ref [ init_env ]
    ; break_lbl = None
    ; continue_lbl = None
    }
  in
  (* 2. 生成线性IR *)
  let raw_code =
    try stmt_to_res ctx0 false (Block f'.body) |> flatten with
    | e ->
      Printf.eprintf
        "Error generating IR for %s: %s\n"
        f'.func_name
        (Printexc.to_string e);
      raise e
  in
  (* 3. 创建入口标签 *)
  let entry_lbl = "entry_" ^ f'.func_name in
  (* 4. 尾调用转换 *)
  let tail_elim_ir = ref [ Label entry_lbl ] in
  List.iter
    (fun inst ->
       match inst with
       | TailCall (fname, args) when fname = f'.func_name ->
         let assigns =
           List.mapi
             (fun i arg ->
                let param = List.nth f'.params i in
                Assign (Var param, arg))
             args
         in
         tail_elim_ir := !tail_elim_ir @ assigns @ [ Goto entry_lbl ]
       | _ -> tail_elim_ir := !tail_elim_ir @ [ inst ])
    raw_code;
  (* 5. 分区和优化 *)
  let raw_blocks = partition_blocks !tail_elim_ir in
  let cfg_blocks = Cfg.build_cfg raw_blocks in
  let opt_blocks = Cfg.optimize cfg_blocks in
  let dce_blocks = dead_code_elimination opt_blocks print_liveness in
  let cse_blocks = List.map (fun x -> cse_block x) dce_blocks in
  let copy_blocks = List.map copy_propagation_block cse_blocks in
  let final_blocks = dead_code_elimination copy_blocks print_liveness in
  (* List.iter dce2_block final_blocks; *)
  { name = f'.func_name; args = f'.params; blocks = final_blocks }
;;

(* 检查表达式是否依赖某些变量 *)
let rec expr_depends_on (e : expr) (vars : StringSet.t) : bool =
  match e with
  | ID x -> StringSet.mem x vars
  | Number _ -> false
  | Binop (_, e1, e2) | Call (_, [ e1; e2 ]) ->
    expr_depends_on e1 vars || expr_depends_on e2 vars
  | Call (_, args) -> List.exists (fun e -> expr_depends_on e vars) args
  | Unop (_, e1) -> expr_depends_on e1 vars
;;

(* 收集当前作用域的赋值目标（写入变量） *)
let rec collect_assigned_vars (s : stmt) : StringSet.t =
  match s with
  | Assign (x, _) -> StringSet.singleton x
  | Decl (x, _) -> StringSet.singleton x
  | Block slist ->
    List.fold_left
      (fun acc s -> StringSet.union acc (collect_assigned_vars s))
      StringSet.empty
      slist
  | If (_, s1, Some s2) ->
    StringSet.union (collect_assigned_vars s1) (collect_assigned_vars s2)
  | If (_, s1, None) -> collect_assigned_vars s1
  | While (_, body) -> collect_assigned_vars body
  | _ -> StringSet.empty
;;

(* 判断一个语句是否循环不变：表达式不依赖于循环中被赋值的变量 *)
let stmt_is_loop_invariant assigned_in_loop stmt =
  match stmt with
  | Decl (_, Some e) -> not (expr_depends_on e assigned_in_loop)
  | Assign (_, e) -> not (expr_depends_on e assigned_in_loop)
  | _ -> false
;;

(* 哪些变量不应外提：被反复初始化，例如 j=0，k=0，t5=0 *)
let collect_reset_vars (stmts : stmt list) : StringSet.t =
  List.fold_left
    (fun acc stmt ->
       match stmt with
       | Decl (x, Some (Number 0)) -> StringSet.add x acc
       | Assign (x, Number 0) -> StringSet.add x acc
       | _ -> acc)
    StringSet.empty
    stmts
;;

(* 遍历语句列表，统计变量被写入的次数 *)
let compute_write_counts (stmts : stmt list) : (string, int) Hashtbl.t =
  let tbl = Hashtbl.create 16 in
  let rec visit (stmt : Ast.stmt) =
    match stmt with
    | Assign (x, _) | Decl (x, _) ->
      let count =
        try Hashtbl.find tbl x with
        | Not_found -> 0
      in
      Hashtbl.replace tbl x (count + 1)
    | Block sl -> List.iter visit sl
    | If (_, s1, Some s2) ->
      visit s1;
      visit s2
    | If (_, s1, None) -> visit s1
    | While (_, body) -> visit body
    | _ -> ()
  in
  List.iter visit stmts;
  tbl
;;

(* 提取循环不变语句，排除被 reset 的变量 *)
let extract_loop_invariants (stmts : stmt list) =
  let assigned_in_loop =
    List.fold_left
      (fun acc s -> StringSet.union acc (collect_assigned_vars s))
      StringSet.empty
      stmts
  in
  let write_counts = compute_write_counts stmts in
  let invariants, rest =
    List.partition
      (fun stmt ->
         let is_reset_target =
           match stmt with
           | Decl (x, Some (Number 0)) | Assign (x, Number 0) | Decl (x, _) | Assign (x, _)
             ->
             (* 多次写入的不允许外提（如 t5） *)
             (try Hashtbl.find write_counts x > 1 with
              | Not_found -> false)
           | _ -> false
         in
         stmt_is_loop_invariant assigned_in_loop stmt && not is_reset_target)
      stmts
  in
  invariants, rest
;;

(* 优化单个语句 *)
let rec optimize_stmt stmt =
  match stmt with
  | While (cond, Block body) ->
    if
      List.exists
        (fun s ->
           match s with
           | Break | Continue -> true
           | _ -> false)
        body
    then While (cond, Block (List.map optimize_stmt body))
    else (
      let body' = List.map optimize_stmt body in
      let invariants, rest = extract_loop_invariants body' in
      Block (invariants @ [ While (cond, Block rest) ]))
  | Block slist -> Block (List.map optimize_stmt slist)
  | If (e, s1, Some s2) -> If (e, optimize_stmt s1, Some (optimize_stmt s2))
  | If (e, s1, None) -> If (e, optimize_stmt s1, None)
  | _ -> stmt
;;

(* 优化函数定义 *)
let optimize_func (f : func_def) : func_def =
  { f with body = List.map optimize_stmt f.body }
;;

(* 优化整个编译单元 *)
let optimize (prog : comp_unit) : comp_unit = List.map optimize_func prog

(* 编译单元转换 *)
let program_to_ir (cu : comp_unit) (optimize_flag : bool) (print_liveness : bool)
  : ir_program
  =
  let cu' = if optimize_flag then cu |> preprocess_ast |> loop_elim_ast |> optimize |> loop_elim_ast else cu in
  (* Printf.printf "\nPrint optimized cu:\n%s\n" (Print_ast.string_of_comp_unit cu'); *)
  let cu' = optimize cu' in
  if optimize_flag
  then Ir_funcs_o (List.map (fun f -> func_to_ir_o f print_liveness) cu')
  else Ir_funcs (List.map func_to_ir cu')
;;
